#include <fcntl.h>
#include <linux/usb/video.h>
#include <linux/uvcvideo.h>
#include <linux/videodev2.h>
#include <ros/ros.h>
#include <sensor_msgs/Image.h>
#include <sensor_msgs/PointCloud2.h>
#include <sensor_msgs/point_cloud2_iterator.h>
#include <string.h>
#include <sys/ioctl.h>
#include <sys/mman.h>

#include <cstdlib>
#include <iostream>

#define BUFFER_COUNT 1
#define FRAME_WIDTH 240
#define FRAME_HEIGHT 180
#define PUB_IR 1
#define PUB_DEPTH 1
#define PUB_POINTCLOUD 1


#define EXTENSION_PROP_ID_CUS 5
static const uint32_t GUID[] = {0x23e49ed0, 0x1178, 0x4f31, 0xae, 0x52, 0xd2, 0xfb, 0x8a, 0x8d, 0x3b, 0x48};

using namespace std;

inline uint32_t UI4Set(uint8_t byte_0, uint8_t byte_1, uint8_t byte_2,
                       uint8_t byte_3) {
  return byte_0 | ((uint32_t)byte_1) << 8 | ((uint32_t)byte_2) << 16 |
         ((uint32_t)byte_3) << 24;
}

inline float FLTSet(uint8_t byte_0, uint8_t byte_1, uint8_t byte_2,
                    uint8_t byte_3) {
  uint32_t flt = UI4Set(byte_0, byte_1, byte_2, byte_3);
  return *reinterpret_cast<float *>(&flt);
}

enum package_type {
  E_INTERGRATION_TIME = 0,
  E_FRAMERATE,
  E_CALIBTRATION_TYPE,
  E_FILTER_TYPE,
  E_AE,
  E_SOC_DEBUG,
  E_HDR_FRAME_TYPE,
};

struct tof_device_intergration_time {
  __uint32_t normal_time;
  __uint32_t hdr_time;
};

struct tof_device_cus_params {
  package_type pkgtype;
  union {
    tof_device_intergration_time inter_time;
    __uint8_t framerate;
    __uint8_t ae;
  };
  __uint8_t direction;
};

struct VideoBuffer {
  void *map_ptr_start;
  size_t length;
};

class JmUVC {
 public:
  JmUVC(const char *device);
  ~JmUVC();
  int getSupportFormat();
  int setUVCFormat();
  int getUVCFormat();
  int getDeviceInfo();
  int videoRequestbufs();
  int videoStreamOn();
  int videoStreamOff();
  int getFrame(float *x, float *y, float *z, uint16_t *ir, float *d);
  int cacheFrame();
  int setIntegration(uint32_t normal_time, uint32_t hdr_time, uint8_t i, uint8_t j);
  int setRate(uint8_t framerate, uint8_t i, uint8_t j);
  int setAE(uint8_t ae, uint8_t i, uint8_t j);

 private:
  int fd_;
  int buffer_num_;
  VideoBuffer frame_buf_[BUFFER_COUNT];
};

JmUVC::JmUVC(const char *device) {
  // 读写打开，拿到设备句柄
  this->fd_ = open(device, O_RDWR);
  if (this->fd_ < 0) {
    ROS_ERROR("%s open failed", device);
    exit(-1);
  } else {
    ROS_DEBUG("%s open success", device);
  }
  getSupportFormat();
  setUVCFormat();
  for(int i=0; i<255; i++) {
    for(int j=0; j<255; j++) {
      setAE(0,i,j);
      setRate(10,i,j);
      setIntegration(1000, 0, i,j);
    }
  }
  
  videoRequestbufs();
  cacheFrame();
  videoStreamOn();
  ros::NodeHandle nh("~");
  ros::Publisher point_pub_ =
      nh.advertise<sensor_msgs::PointCloud2>("/jm/point", 1);
  ros::Publisher depth_pub_ = nh.advertise<sensor_msgs::Image>("/jm/depth", 1);
  ros::Publisher ir_pub_ = nh.advertise<sensor_msgs::Image>("/jm/ir", 1);
  float *x = new float[FRAME_WIDTH * FRAME_HEIGHT];
  float *y = new float[FRAME_WIDTH * FRAME_HEIGHT];
  float *z = new float[FRAME_WIDTH * FRAME_HEIGHT];
  uint16_t *ir = new uint16_t[FRAME_WIDTH * FRAME_HEIGHT];
  float *depth = new float[FRAME_WIDTH * FRAME_HEIGHT];
  while (ros::ok()) {
    ros::Time now_t = ros::Time::now();
    getFrame(x, y, z, ir, depth);
    cacheFrame();

#if PUB_POINTCLOUD
    sensor_msgs::PointCloud2 pc;  
    pc.header.frame_id = "jm";
    pc.header.stamp = now_t;
    pc.height = 1;
    pc.width = FRAME_WIDTH * FRAME_HEIGHT;
    sensor_msgs::PointCloud2Modifier modifier(pc);
    modifier.setPointCloud2Fields(3, "x", 1, sensor_msgs::PointField::FLOAT32,
                                  "y", 1, sensor_msgs::PointField::FLOAT32, "z",
                                  1, sensor_msgs::PointField::FLOAT32);
    sensor_msgs::PointCloud2Iterator<float> iter_x(pc, "x");
    sensor_msgs::PointCloud2Iterator<float> iter_y(pc, "y");
    sensor_msgs::PointCloud2Iterator<float> iter_z(pc, "z");

    for (int i = 0; i < FRAME_WIDTH; i++) {
      for (int j = 0; j < FRAME_HEIGHT; j++) {
        uint16_t ir_v = *(ir + j * FRAME_WIDTH + i);
        if(ir_v > 1000 || ir_v < 50) {
          *iter_x = 0;
          *iter_y = 0;
          *iter_z = 0;
        } else {
          *iter_x = *(x + j * FRAME_WIDTH + i);
          *iter_y = *(y + j * FRAME_WIDTH + i);
          *iter_z = *(z + j * FRAME_WIDTH + i);
        }
        ++iter_x;
        ++iter_y;
        ++iter_z;
      }
    }
    point_pub_.publish(pc);
#endif

#if PUB_IR
    sensor_msgs::Image ir_img;
    ir_img.header.stamp = now_t;
    ir_img.header.frame_id = "jm";
    ir_img.width = FRAME_WIDTH;
    ir_img.height = FRAME_HEIGHT;
    ir_img.encoding = "16SC1";
    ir_img.is_bigendian = false;
    ir_img.step = 2 * FRAME_WIDTH;
    ir_img.data.resize(2 * FRAME_WIDTH * FRAME_HEIGHT);
    for (int i = 0; i < FRAME_WIDTH; i++) {
      for (int j = 0; j < FRAME_HEIGHT; j++) {
        uint16_t v = *(ir + j * FRAME_WIDTH + i);
        ir_img.data[2 * (FRAME_HEIGHT - 1 - j) * FRAME_WIDTH +
                    2 * (FRAME_WIDTH - 1 - i)] = (uint8_t)v;
        ir_img.data[2 * (FRAME_HEIGHT - 1 - j) * FRAME_WIDTH +
                    2 * (FRAME_WIDTH - 1 - i) + 1] = v >> 8;
      }
    }
    ir_pub_.publish(ir_img);
#endif

#if PUB_DEPTH
    sensor_msgs::Image depth_img;
    depth_img.header.stamp = now_t;
    depth_img.header.frame_id = "jm";
    depth_img.width = FRAME_WIDTH;
    depth_img.height = FRAME_HEIGHT;
    depth_img.encoding = "32FC1";
    depth_img.is_bigendian = false;
    depth_img.step = 4 * FRAME_WIDTH;
    depth_img.data.resize(4 * FRAME_WIDTH * FRAME_HEIGHT);
    for (int i = 0; i < FRAME_WIDTH; i++) {
      for (int j = 0; j < FRAME_HEIGHT; j++) {
        float v = *(depth + j * FRAME_WIDTH + i);
        uint8_t *arr;
        arr = reinterpret_cast<uint8_t *>(&v);
        depth_img.data[4 * (FRAME_HEIGHT - 1 - j) * FRAME_WIDTH +
                       4 * (FRAME_WIDTH - 1 - i)] = *arr;
        depth_img.data[4 * (FRAME_HEIGHT - 1 - j) * FRAME_WIDTH +
                       4 * (FRAME_WIDTH - 1 - i) + 1] = *(arr + 1);
        depth_img.data[4 * (FRAME_HEIGHT - 1 - j) * FRAME_WIDTH +
                       4 * (FRAME_WIDTH - 1 - i) + 2] = *(arr + 2);
        depth_img.data[4 * (FRAME_HEIGHT - 1 - j) * FRAME_WIDTH +
                       4 * (FRAME_WIDTH - 1 - i) + 3] = *(arr + 3);
      }
    }
    depth_pub_.publish(depth_img);
#endif
  }

  videoStreamOff();
}

JmUVC::~JmUVC() {}

int JmUVC::getDeviceInfo() {
  int retv;
  v4l2_capability cap;
  // 获取有关驱动程序和硬件功能的信息。
  retv = ioctl(fd_, VIDIOC_QUERYCAP, &cap);
  if (retv < 0) {
    ROS_ERROR("get device infomation failed");
  } else {
    ROS_DEBUG("get device infomation success");
  }
  return retv;
}

int JmUVC::getSupportFormat() {
  int retv;
  v4l2_fmtdesc fmt;
  memset(&fmt, 0, sizeof(fmt));
  fmt.index = 0;
  fmt.type = V4L2_BUF_TYPE_VIDEO_CAPTURE;
  while ((retv = ioctl(fd_, VIDIOC_ENUM_FMT, &fmt)) == 0) {
    fmt.index++;
    printf("{ pixelformat = ''%c%c%c%c'', description = ''%s'' }\n",
           fmt.pixelformat & 0xFF, (fmt.pixelformat >> 8) & 0xFF,
           (fmt.pixelformat >> 16) & 0xFF, (fmt.pixelformat >> 24) & 0xFF,
           fmt.description);
  }
  return retv;
}

int JmUVC::getUVCFormat() {
  int retv;
  v4l2_format fmt;
  memset(&fmt, 0, sizeof(fmt));
  fmt.type = V4L2_BUF_TYPE_VIDEO_CAPTURE;
  retv = ioctl(this->fd_, VIDIOC_G_FMT, &fmt);
  if (retv < 0) {
    ROS_ERROR("getUVCFormat failed");
  } else {
    ROS_DEBUG("getUVCFormat success");
    ROS_INFO("width = %d, height = %d", fmt.fmt.pix.width, fmt.fmt.pix.height);
    printf("{ pixelformat = ''%c%c%c%c'' }\n", fmt.fmt.pix.pixelformat & 0xFF,
           (fmt.fmt.pix.pixelformat >> 8) & 0xFF,
           (fmt.fmt.pix.pixelformat >> 16) & 0xFF,
           (fmt.fmt.pix.pixelformat >> 24) & 0xFF);
  }
  return retv;
}

int JmUVC::setUVCFormat() {
  int retv;
  v4l2_format fmt;
  memset(&fmt, 0, sizeof(fmt));
  fmt.type = V4L2_BUF_TYPE_VIDEO_CAPTURE;
  fmt.fmt.pix.width = 3840;
  fmt.fmt.pix.height = 2160;
  fmt.fmt.pix.pixelformat = V4L2_PIX_FMT_H264;
  retv = ioctl(this->fd_, VIDIOC_S_FMT, &fmt);
  if (retv < 0) {
    ROS_ERROR("setUVCFormat failed");
  } else {
    ROS_DEBUG("setUVCFormat success");
    getUVCFormat();
  }
  return retv;
}

int JmUVC::videoRequestbufs() {
  int retv;
  v4l2_requestbuffers rqBuf;
  rqBuf.type = V4L2_BUF_TYPE_VIDEO_CAPTURE;
  rqBuf.count = BUFFER_COUNT;
  rqBuf.memory = V4L2_MEMORY_MMAP;
  retv = ioctl(fd_, VIDIOC_REQBUFS, &rqBuf);
  if (retv < 0) {
    ROS_ERROR("videoRequestbufs failed");
    return retv;
  } else {
    ROS_DEBUG("videoRequestbufs success");
    v4l2_buffer buf;
    for (unsigned int i = 0; i < rqBuf.count; i++) {
      memset(&buf, 0, sizeof(buf));
      buf.index = i;
      buf.type = V4L2_BUF_TYPE_VIDEO_CAPTURE;
      buf.memory = V4L2_MEMORY_MMAP;
      retv = ioctl(fd_, VIDIOC_QUERYBUF, &buf);
      if (retv < 0) {
        ROS_ERROR("query buffer failed");
        return retv;
      } else {
        ROS_DEBUG("query buffer success");
      }
      frame_buf_[i].length = buf.length;
      frame_buf_[i].map_ptr_start = mmap(0, buf.length, PROT_READ | PROT_WRITE,
                                         MAP_SHARED, fd_, buf.m.offset);
      if (frame_buf_[i].map_ptr_start == MAP_FAILED) {
        ROS_ERROR("mmap[%d] failed", i);
        return -1;
      } else {
        ROS_DEBUG("mmap[%d] success", i);
      }
    }
    buffer_num_ = rqBuf.count;
  }
  return retv;
}

int JmUVC::cacheFrame() {
  int retv;
  v4l2_buffer buf;
  for (unsigned int i = 0; i < buffer_num_; i++) {
    buf.index = i;
    buf.type = V4L2_BUF_TYPE_VIDEO_CAPTURE;
    buf.memory = V4L2_MEMORY_MMAP;
    retv = ioctl(fd_, VIDIOC_QBUF, &buf);
    if (retv < 0) {
      ROS_ERROR("cacheFrame[%d] failed", i);
      return retv;
    } else {
      ROS_DEBUG("cacheFrame[%d] succes", i);
    }
  }
  return retv;
}

int JmUVC::getFrame(float *x, float *y, float *z, uint16_t *ir, float *d) {
  int retv;
  v4l2_buffer buf;
  for (unsigned int i = 0; i < buffer_num_; i++) {
    buf.index = i;
    buf.type = V4L2_BUF_TYPE_VIDEO_CAPTURE;
    buf.memory = V4L2_MEMORY_MMAP;
    retv = ioctl(fd_, VIDIOC_DQBUF, &buf);
    if (retv < 0) {
      ROS_ERROR("getFrame[%d] failed", i);
      return retv;
    } else {
      uint32_t frame_length = FRAME_WIDTH * FRAME_HEIGHT;
      uint32_t start = 0;
      uint32_t end = frame_length * sizeof(uint16_t);
      uint8_t *content = (uint8_t *)frame_buf_[buf.index].map_ptr_start;
      int ind = 0;
// #if PUB_IR
      for (uint32_t j = start; j < end;) {
        uint8_t ir1 = *(content + j);
        uint8_t ir2 = *(content + ++j);
        j++;
        *(ir + ind) = ir2 << 8 | ir1;
        ind++;
      }
// #endif

      start = end;
      end = start + frame_length * sizeof(float);
      ind = 0;
#if PUB_DEPTH
      for (uint32_t j = start; j < end;) {
        uint8_t d1 = *(content + j);
        uint8_t d2 = *(content + ++j);
        uint8_t d3 = *(content + ++j);
        uint8_t d4 = *(content + ++j);
        j++;
        *(d + ind) = FLTSet(d1, d2, d3, d4);
        ind++;
      }
#endif

      ind = 0;
      start = end;
      end = start + frame_length * sizeof(float) * 3;
#if PUB_POINTCLOUD
      for (uint32_t j = start; j < end;) {
        uint8_t xb1 = *(content + j);
        uint8_t xb2 = *(content + ++j);
        uint8_t xb3 = *(content + ++j);
        uint8_t xb4 = *(content + ++j);
        uint8_t yb1 = *(content + ++j);
        uint8_t yb2 = *(content + ++j);
        uint8_t yb3 = *(content + ++j);
        uint8_t yb4 = *(content + ++j);
        uint8_t zb1 = *(content + ++j);
        uint8_t zb2 = *(content + ++j);
        uint8_t zb3 = *(content + ++j);
        uint8_t zb4 = *(content + ++j);
        j++;
        *(x + ind) = FLTSet(xb1, xb2, xb3, xb4);
        *(y + ind) = FLTSet(yb1, yb2, yb3, yb4);
        *(z + ind) = FLTSet(zb1, zb2, zb3, zb4);
        ind++;
      }
#endif
    }
  }
  return retv;
}

int JmUVC::videoStreamOn() {
  int retv;
  v4l2_buf_type v4l2type = V4L2_BUF_TYPE_VIDEO_CAPTURE;
  retv = ioctl(fd_, VIDIOC_STREAMON, &v4l2type);
  if (retv < 0) {
    ROS_ERROR("videoStreamOn failed");
  } else {
    ROS_DEBUG("videoStreamOn failed");
  }
  return retv;
}

int JmUVC::videoStreamOff() {
  int retv;
  v4l2_buf_type v4l2type = V4L2_BUF_TYPE_VIDEO_CAPTURE;
  retv = ioctl(fd_, VIDIOC_STREAMOFF, &v4l2type);
  if (retv < 0) {
    ROS_ERROR("videoStreamOff failed");
  } else {
    ROS_DEBUG("videoStreamOff failed");
    for(int i = 0; i < BUFFER_COUNT; i++) {
      munmap(frame_buf_[i].map_ptr_start, frame_buf_[i].length);
    }
  }
  return retv;
}

int JmUVC::setIntegration(uint32_t normal_time, uint32_t hdr_time, uint8_t i, uint8_t j) {
  int retv;
  tof_device_cus_params cus_config;
  cus_config.pkgtype = E_INTERGRATION_TIME;
  cus_config.inter_time.normal_time = normal_time;
  // cus_config.inter_time.hdr_time = hdr_time;
  uvc_xu_control_query xu_ctrl_query;
  // xu_ctrl_query.unit = EXTENSION_PROP_ID_CUS;  
  xu_ctrl_query.unit = i;
  xu_ctrl_query.selector = j;
  xu_ctrl_query.query = UVC_SET_CUR;
  xu_ctrl_query.size = sizeof(cus_config);
  xu_ctrl_query.data =  (unsigned char*)&cus_config;
  retv = ioctl(this->fd_, UVCIOC_CTRL_QUERY, xu_ctrl_query);
  if (retv < 0) {
    // ROS_ERROR("setIntegration failed");
  } else {
    ROS_DEBUG("setIntegration success");
    ROS_INFO("%d, %d", i, j);
  }
  return retv;
}

int JmUVC::setRate(uint8_t framerate, uint8_t i, uint8_t j) {
  int retv;
  tof_device_cus_params cus_config;
  cus_config.pkgtype = E_FRAMERATE;
  cus_config.framerate = framerate;
  uvc_xu_control_query xu_ctrl_query;
  // xu_ctrl_query.unit = EXTENSION_PROP_ID_CUS;
  xu_ctrl_query.unit = i;
  xu_ctrl_query.selector = j;
  xu_ctrl_query.query = UVC_SET_CUR;
  xu_ctrl_query.size = sizeof(cus_config);
  xu_ctrl_query.data = (unsigned char*)&cus_config;
  retv = ioctl(this->fd_, UVCIOC_CTRL_QUERY, xu_ctrl_query);
  if (retv < 0) {
    // ROS_ERROR("setRate failed");
  } else {
    ROS_DEBUG("setRate success");
    ROS_INFO("%d, %d", i, j);
  }
  return retv;
}

int JmUVC::setAE(uint8_t ae, uint8_t i, uint8_t j) {
  int retv;
  tof_device_cus_params cus_config;
  cus_config.pkgtype = E_AE;
  cus_config.ae = ae;
  uvc_xu_control_query xu_ctrl_query;
  // xu_ctrl_query.unit = EXTENSION_PROP_ID_CUS;
  xu_ctrl_query.unit = i;
  xu_ctrl_query.selector = j;
  xu_ctrl_query.query = UVC_SET_CUR;
  xu_ctrl_query.size = sizeof(cus_config);
  xu_ctrl_query.data = (unsigned char*)&cus_config;
  retv = ioctl(this->fd_, UVCIOC_CTRL_QUERY, xu_ctrl_query);
  if (retv < 0) {
    // ROS_ERROR("setAE failed");
  } else {
    ROS_DEBUG("setAE success");
    ROS_INFO("%d, %d", i, j);
  }
  return retv;
}

int main(int argc, char *argv[]) {
  ros::init(argc, argv, "jm_uvc");
  char *port = argv[1];
  JmUVC jmUVC(port);
  return 0;
}